import os
import SimpleITK as sitk
import pydicom
import numpy as np


def get_root_dcm_files(root_start: str):
    root_dcm_dict = {}

    for root, dirs, files in os.walk(root_start):
        filter_dcm_file = []
        for each_file in files:
            if each_file.lower().endswith(".dcm"):
                filter_dcm_file.append(each_file)

        if len(filter_dcm_file) == 0:
            continue

        # TODO: 文件夹可能存在多个脏数据
        reader = sitk.ImageSeriesReader()
        root_dcm_dict[root] = reader.GetGDCMSeriesIDs(root)[-1]
        print("Read", root, root_dcm_dict[root])

    return root_dcm_dict


def convert_all(root_dcm_dict: dict, save_path: str):
    os.makedirs(save_path, exist_ok=True)
    for root in root_dcm_dict.keys():
        try:
            reader = sitk.ImageSeriesReader()
            series_file_names = reader.GetGDCMSeriesFileNames(root, root_dcm_dict[root])

            ds = pydicom.dcmread(series_file_names[0])
            save_name = ds.AccessionNumber + ".nrrd"

            reader.SetFileNames(series_file_names)
            dicom_series = reader.Execute()

            sitk.WriteImage(dicom_series, os.path.join(save_path, save_name))
            print("Convert", save_name)
        except Exception as e:
            print("Convert Error", root_dcm_dict[root], str(e))


def convert_all_v2(root_dcm_dict: dict, save_path: str):
    os.makedirs(save_path, exist_ok=True)
    for root in root_dcm_dict.keys():
        try:
            reader = sitk.ImageSeriesReader()
            series_file_names = reader.GetGDCMSeriesFileNames(root, root_dcm_dict[root])

            ds = pydicom.dcmread(series_file_names[0])
            save_name = ds.AccessionNumber + ".nrrd"
            spacing = ds.PixelSpacing
            spacing.append(ds.SliceThickness)
            orientation = ds.ImageOrientationPatient
            position = ds.ImagePositionPatient

            reader.SetFileNames(series_file_names)
            dicom_series = reader.Execute()

            dicom_series.SetSpacing(spacing)
            direction_cosines = np.array(orientation).reshape(2, 3)
            normal = np.cross(direction_cosines[0], direction_cosines[1])
            direction_cosines = np.vstack([direction_cosines, normal])  # 添加法向量
            dicom_series.SetDirection(direction_cosines.flatten())
            dicom_series.SetOrigin(position)

            sitk.WriteImage(dicom_series, os.path.join(save_path, save_name))
            print("Convert", save_name)
        except Exception as e:
            print("Convert Error", root_dcm_dict[root], str(e))


if __name__ == '__main__':
    start = r"Z:\Development\Ftp\dicom_data_test"
    end = "Z:\Development\Ftp\dicom_save"
    d = get_root_dcm_files(start)
    convert_all(d, end)